import os

from matplotlib import pyplot as plt
from PIL.Image import fromarray

from Iterator import random_walk, FaceDataset


# 图片目录
pl0 = os.listdir("./Resource")[0:10]
path_list = ["./Resource/"+p for p in pl0]

# 测试随机游走生成器，生成5个数据
print("Test 1\n随机游走生成器要求1\n")
rw1 = random_walk(1,1,1,3)
for i in range(5):
    print(next(rw1))
print("\nTest 1 Complete.\n")

# 测试拼合而成的随机游走生成器，可以生成两组对齐的多维随机游走序列
print("Test 2\n随机游走生成器要求2\n")
rw2 = random_walk(1,2,1,3)
rw3 = random_walk(1,2,1,4)
u = zip(rw2, rw3)
for i in range(3):
    print(next(u))
print("\nTest 2 Complete.\n")

# 测试FaceDataset的__next__方法：路径列表长度为10，运行11次next函数
# 应当输出10张图片，以及一个提示信息。
print("Test 3\nFaceDataset __next__方法测试")
fds = FaceDataset(path_list)
print(f"面部数据集大小为{len(fds)}")
try:
    for i in range(11):
        next(fds)
        print(f"已加载{len(fds.images)}张图片")
        plt.figure("Test Display")
        plt.imshow(fromarray(fds.images[-1]))
        plt.axis('off')
        plt.title(fds.path_list[i])
        plt.show()
except StopIteration as s:
    print(*s.args)
print("\nTest 3 Complete.\n")

# 测试FaceDataset的__iter__方法：应当从头加载10张图片。
print("Test 4\nFaceDataset __iter__方法测试")
for face in fds:
    print(f"已加载{len(fds.images)}张图片")
print("\nTest 4 Complete.\n")
